import torch
import torchvision
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor
from torch.optim import SGD
import time

import MyModel

train_dataset = CIFAR10("./dataset",
                         transform=ToTensor(),
                         train=True,
                         download=True)
test_dataset = CIFAR10("./dataset",
                         transform=ToTensor(),
                         train=False,
                         download=True)
# 训练集和测试集的大小
len_train_dataset = len(train_dataset)
len_test_dataset = len(test_dataset)
print(f"训练集的数据量大小为:{len_train_dataset}")
print(f"训练集的数据量大小为:{len_test_dataset}")

# 训练的batch数据集
train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=True)

# 学习率，epoch
learning_rate = 1e-2
epoch = 10

# 创建模型，损失函数，优化器
my_model = MyModel.CIFAR10Model()
loss_func = CrossEntropyLoss()
optimizer = SGD(my_model.parameters(), lr=learning_rate)

# 轮数
count = 0

# 数据记入tensorboard
writer = SummaryWriter('./log')

# 训练过程
for i in range(epoch):
    start_time = time.time()
    for images, labels in train_dataloader:
        optimizer.zero_grad()
        outputs = my_model(images)
        loss = loss_func(outputs, labels)
        loss.backward()
        optimizer.step()

        # 打印第几轮学习
        if count % 100 == 0:
            print(f"----------第{i + 1}轮训练, 第{count}轮batch, 损失值:{loss.item()}------------")
            writer.add_scalar('train_loss', loss.item(), count)
        count += 1

    print(f"----------第{i+1}轮epoch训练用时：{time.time() - start_time :.2f}s--------")
    # 每个epoch训练完成后保存模型
    torch.save(my_model, f"my_model_{i+1}.pth" )

    correct = 0
    test_loss = 0
    # 测试集准确率
    with torch.no_grad():
        for images, labels in test_dataloader:
            outputs = my_model(images)
            correct += (outputs.argmax(1) == labels).sum()

            test_loss += loss_func(outputs, labels).item()

    print(f"测试集总损失值：{test_loss}, 准确率：{correct/len_test_dataset*100 :.2f}%")
    writer.add_scalar('test_lossasdfasdf', test_loss, i)


writer.close()